# main.py
from generate_loader import generate_loader
from encrypt_loader import encrypt_loader
import subprocess
import sys
import os

def pack_py2exe(py_file: str):
    setup_py = "setup_py2exe.py"
    # 生成简单的 setup 脚本
    with open(setup_py, "w", encoding="utf-8") as f:
        f.write(f"""
from distutils.core import setup
import py2exe
import sys
sys.argv.append('py2exe')

setup(
    options={{"py2exe": {{"bundle_files": 1, "compressed": True, "optimize": 2}}}},
    windows=[{{"script": "{py_file}"}}],
    zipfile=None,
)
""")

    print("[*] 使用 py2exe 打包...")
    result = subprocess.run([sys.executable, setup_py], capture_output=True, text=True)
    if result.returncode == 0:
        print("[+] py2exe 打包成功，生成的 exe 在 dist/ 目录下")
    else:
        print("[!] py2exe 打包失败:")
        print(result.stdout)
        print(result.stderr)

    # 删除临时 setup 文件（如果需要）
    if os.path.exists(setup_py):
        os.remove(setup_py)

def pack_to_exe(py_file: str, exe_name: str, onefile: bool = True, noconsole: bool = False) -> bool:
    cmd = [sys.executable, "-m", "PyInstaller"]

    if onefile:
        cmd.append("--onefile")
    if noconsole:
        cmd.append("--noconsole")

    # 显式指定必须导入的模块
    hidden_imports = [
        "Cryptodome"
        # "Cryptodome.Cipher.AES",
        # "Cryptodome.Cipher.DES",  # 如果你使用了 DESPlugin
        # "Cryptodome.Util.Padding",
        # "encrypt_plugins.aes_plugin",
        # "encrypt_plugins.xor_plugin",
        # "encrypt_plugins.base64_plugin",
        # "encrypt_plugins.custom_base64_plugin",
        # "templates.exec_template"
    ]
    for module in hidden_imports:
        cmd.append(f"--hidden-import={module}")

    cmd.extend(["--name", exe_name, py_file])

    print(f"[*] 运行命令：{' '.join(cmd)}")
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        print(f"[+] PyInstaller 打包成功，生成的 exe 在 dist/{exe_name}.exe")
        return True
    else:
        print("[!] PyInstaller 打包失败:")
        print(result.stdout)
        print(result.stderr)
        return False


def pack_cxfreeze(script_file):
    setup_script = "setup_cxfreeze.py"
    setup_code = f"""
import sys
from cx_Freeze import setup, Executable

base = None
if sys.platform == "win32":
    pass

setup(
    name="loader_enc",
    version="1.0",
    description="Encrypted loader executable",
    options={{"build_exe": {{
        "packages": ["Cryptodome", "encrypt_plugins", "templates"],
        "includes": ["encrypt_plugins.aes_plugin", "encrypt_plugins.xor_plugin", "encrypt_plugins.base64_plugin", "encrypt_plugins.custom_base64_plugin", "templates.exec_template"]
    }}}},
    executables=[Executable("{script_file}", base=base)]
)
"""
    with open(setup_script, "w", encoding="utf-8") as f:
        f.write(setup_code)

    print("[*] 使用 cx_Freeze 打包...")
    result = subprocess.run([sys.executable, setup_script, "build"], capture_output=True, text=True)
    if result.returncode == 0:
        print("[+] cx_Freeze 打包成功，生成的 exe 在 build 目录下")
    else:
        print("[!] cx_Freeze 打包失败:")
        print(result.stdout)
        print(result.stderr)

    # if os.path.exists(setup_script):
    #     os.remove(setup_script)

def pack_nuitka(py_file: str, exe_name: str):
    print("[*] 使用 Nuitka 打包...")
    cmd = [
        sys.executable, "-m", "nuitka",
        "--onefile",
        f"--output-filename={exe_name}.exe",
        "--include-package=Cryptodome",
        "--include-package=encrypt_plugins",
        "--include-package=templates",
        py_file
    ]

    print(f"[*] 运行命令：{' '.join(cmd)}")
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        print(f"[+] Nuitka 打包成功，生成的 exe 在当前目录: {exe_name}.exe")
    else:
        print("[!] Nuitka 打包失败:")
        print(result.stdout)
        print(result.stderr)


def main():
    print("=== 生成 loader.py ===")
    # generate_loader()

    print("\n=== 对 loader.py 进行加密生成 loader_enc.py ===")
    # encrypt_loader()

    print("\n=== 选择打包方式 ===")
    print("1. py2exe")
    print("2. cx_Freeze")
    print("3. PyInstaller")
    print("4. Nuitka")
    choice = input("请输入序号 (1/2/3/4), 或其他退出: ").strip()

    if choice == "1":
        pack_py2exe("loader_enc.py")
    elif choice == "2":
        pack_cxfreeze("loader_enc.py")
    elif choice == "3":
        success = pack_to_exe("loader_enc.py", "loader_enc", onefile=True, noconsole=False)
        if not success:
            print("[!] PyInstaller 打包失败，请检查错误信息。")
    elif choice == "4":
        pack_nuitka("loader_enc.py", "loader_enc")
    else:
        print("退出程序，不打包。")

    print("\n[+] 全部流程完成！")

if __name__ == "__main__":
    main()
